#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2018/12/30 0030 15:13
# @Author : Ares_Wang
# @Site : 
# @File : other_model_train.py
# @Software: PyCharm
from src.utils import *
from keras.optimizers import SGD
from src.model.copy.GoogleNet import GoogleNet
from src.model.copy.AlexNet import AlexNet
import numpy as np

# 测试网络
if __name__ == '__main__':

    epochs = 50
    train_filename = r'H:\wangjianlian\data\covered_imgs_test\train\20X'
    batch_size = 100
    if_save_model = True

    # googleNet = GoogleNet()
    # model = googleNet.build()

    alexNet = AlexNet()
    model = alexNet.build()

    model.compile(optimizer=SGD(lr=0.0005, momentum=0.9, nesterov=True),
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])

    data = Data()
    max_metrics = 0
    for epoch in range(epochs):
        print("************ 第%d代 *************" % (epoch + 1))
        all_loss = 0
        all_metrics = 0
        for index, trained_imgs, trained_label in data.read_image(train_filename, batch_size, img_width = 227, img_height = 227, shuffle=True):
            x_train = trained_imgs
            y_train = trained_label[0]

            # y_train = np.expand_dims(y_train, axis=1)  # 增维
            # y_train = np.expand_dims(y_train, axis=1)

            loss_and_metrics = model.train_on_batch(x_train, y_train)  # 应该是返回损失值和metrics
            print("训练第%d批，loss值:%f" % (index, loss_and_metrics[0]))
            print("训练第%d批，metrics值:%f" % (index, loss_and_metrics[1]))
            all_loss += loss_and_metrics[0]
            all_metrics += loss_and_metrics[1]

        print("************ 训练：平均值 *************")
        print("训练平均loss值:%f" % (all_loss / index))
        print("训练平均metrics值:%f" % (all_metrics / index))

        if if_save_model:
            if all_metrics / index > max_metrics:
                max_metrics = all_metrics / index
                model.save(r"H:\wangjianlian\project\Python\networkTest\resources\weight\temp\AlexNet" + "//" + str(
                    epoch + 1) + 'model_1.h5')
    # model.fit(data, labels)